import numpy as np
import torch
import matplotlib.pyplot as plt
from mmcls.models import build_classifier
from mmcv.runner import load_checkpoint
from mmcv import Config
from mmcls.datasets import build_dataloader, build_dataset

# 定义获取损失表面的函数
def get_loss_surface(model, dataloader, direction1, direction2, alpha_range, beta_range):
    original_params = [p.clone() for p in model.parameters()]
    losses = []

    for alpha in alpha_range:
        row = []
        for beta in beta_range:
            # 在参数空间移动
            with torch.no_grad():
                for p, d1, d2 in zip(model.parameters(), direction1, direction2):
                    p.copy_(p + alpha * d1 + beta * d2)

            # 计算损失
            model.eval()
            total_loss = 0
            for batch in dataloader:
                inputs = batch['img'].to(next(model.parameters()).device)
                targets = batch['gt_label'].to(next(model.parameters()).device)
                outputs = model(inputs)
                loss = torch.nn.CrossEntropyLoss()(outputs, targets)
                total_loss += loss.item()

            row.append(total_loss / len(dataloader))
        
        losses.append(row)

    # 恢复原始参数
    with torch.no_grad():
        for p, original_p in zip(model.parameters(), original_params):
            p.copy_(original_p)

    return np.array(losses)

# 配置模型和数据
config_path = 'path_to_config.py'
checkpoint_path = 'path_to_checkpoint.pth'
cfg = Config.fromfile(config_path)

# 构建模型
model = build_classifier(cfg.model)
load_checkpoint(model, checkpoint_path, map_location='cpu')
model = model.cuda()

# 构建数据加载器
dataset = build_dataset(cfg.data.test, default_args=dict(test_mode=True))
dataloader = build_dataloader(dataset, batch_size=32, shuffle=False, num_workers=4)

# 生成随机方向
direction1 = [torch.randn_like(p) for p in model.parameters()]
direction2 = [torch.randn_like(p) for p in model.parameters()]

# 归一化方向
direction1 = [d / d.norm() for d in direction1]
direction2 = [d / d.norm() for d in direction2]

# 定义平面范围
alpha_range = np.linspace(-1, 1, 50)
beta_range = np.linspace(-1, 1, 50)

# 计算损失表面
loss_surface = get_loss_surface(model, dataloader, direction1, direction2, alpha_range, beta_range)

# 可视化
alpha_grid, beta_grid = np.meshgrid(alpha_range, beta_range)
fig = plt.figure(figsize=(10, 8))
ax = fig.add_subplot(111, projection='3d')
ax.plot_surface(alpha_grid, beta_grid, loss_surface, cmap='viridis')
ax.set_xlabel('Alpha')
ax.set_ylabel('Beta')
ax.set_zlabel('Loss')
plt.savefig('loss_surface_vit.png')
